using Pkg
Pkg.activate(".")

using ComponentArrays, Lux, DiffEqFlux, OrdinaryDiffEq, Optimization, OptimizationOptimJL, ReverseDiff,
    OptimizationOptimisers, Random, Plots, Statistics, Zygote, JLD2, Distributed, Setfield, LinearAlgebra
include("agentzoo.jl")
using .agentzoo
rng = Random.default_rng()
mydist = false;#是否并行
if mydist
    addprocs(1) # 添加额外的处理器
    all_procs = procs()
    println("所有进程的ID：", all_procs)# 输出所有进程的ID
    myEnsemble = EnsembleDistributed()#允许并行
else
    myEnsemble = EnsembleThreads()#只是多线程，不并行
end

#实用的函数
function random_grouping(N, group_size)
    #把1-N的数打乱分成多组，每组数量不超过group_size
    sequence = 1:N

    # 打乱序列
    sequence = Random.shuffle(sequence)

    # 计算组数
    group_num = cld(N, group_size)  # cld是向上取整函数

    # 分成四组
    groups = [sequence[i:i+group_size-1] for i in 1:group_size:group_size*group_num]
    return groups
end

#########
tsize = 3
tspan = (0.0, 1500.0)
ndt = 1e-2;#NODE计算步长
tsteps = range(tspan[1], tspan[2]; length=tsize)
tstepsdata = range(tspan[1], tspan[2]; length=100)
expochs = 1000;#外部训练次数
cpochs = 1;#保存的间隔数量
inpochs = 1;#每个batch训练次数
BFGSpoch = 2;#使用其他优化器的间隔
inBFGSpochs = 1;#每次BFGS训练次数
lr = 1e-4;#学习率
sampnum = 2048;#样本弹道数量
batchnum = 64;#每次训练时抽样的弹道数量
ngreedy = 0.01;#归一化量的软更新比
stdeps = 1e-4;#归一化量的最小值
sampsign = true;#1为重新生成弹道，0为加载弹道
statenum = 8;#状态量数量
goalnum = 6;#指示目标的变量数量
pathnum = 0;#用于记录过程量的变量数量
actnum = 2;#网络输出的动作数量

maxang = 0.02;#动作幅度
#########规划约束量
Verrormax = 10.0;#终端速度误差最大值
Sferrormax = 1e3;#剩余航程误差最大值
rerrormax = 1e3;#终端高度最大值
#########动力学部分需要的参数和函数
R0 = 6378135;
g0 = 9.80665;
OMEGA = 7.292115e-5;
t_scale = sqrt(R0 / g0);
V_scale = sqrt(g0 * R0);

Pa0_atmos = 1.01325e5;
hs_atmos = 7110;
rho0 = 1.225;
kesim = 2.26e-3;

miu = 3.986005e14;
pi = 3.141592653589793;
DTR = pi / 180.0;
RTD = 180.0 / pi;
#热流计算相关
kq = 34879.0;
#限幅相关
magsigma = 80.0 * pi / 180;#倾侧角最大幅度
magalpha = 20.0 * pi / 180;#攻角最大幅度
minalpha = 4.0 * pi / 180;#攻角最小值
#过程约束相关
maxq = 200000;#过程约束，动压约束
maxdQ = 1200;#过程约束，热流约束
maxnn = 4 * g0;#过程约束，过载约束
#气动拟合公式系数
ad = 0.16172;
bd = -0.018452;
cd = 0.001565;
dd = 0.000709;
ed = -0.000382;
fd = 0.000830;

al = -0.047006;
bl = -0.006756;
cl = 0.050686;
dl = 0.000416;
el = -0.000875;
fl = 0.000374;

mass = 907.0;
Sref = 0.35;
function mybound(x)
    y = asin.(sin.(x)) ./ (pi / 2) .* maxang
    return y
end
function sigmoid_clamp(x, low, high)
    scale = high - low
    mid = (high + low) / 2
    return scale ./ (1 .+ exp.(-x)) .+ mid
end
# function m_GetDensity(href)
#     #大气密度是高度的指数函数
#     #h单位是km
#     R = 6356.766
#     rho0 = 1.225
#     href=clamp(href, 0.0, 120.0)
#     hp = R * href / (R + href)
#     if href <= 11.0191
#         W = 1 - (hp / 44.3308)
#         rho = rho0 * W^4.2559
#         T = 288.15 * W
#         Pa = Pa0_atmos * W^5.2559
#     elseif 11.0191 < href <= 20.0631
#         W = exp((14.9647 - hp) / 6.3416)
#         rho = 0.15898 * rho0 * W
#         T = 216.650
#         Pa = Pa0_atmos * 0.11953 * W
#     elseif 20.0631 < href <= 32.1619
#         W = 1 + ((hp - 24.9021) / 221.552)
#         rho = 0.032722 * rho0 * W^(-35.1629)
#         T = 221.552 * W
#         Pa = 0.025158 * Pa0_atmos * W^(-34.1629)
#     elseif 32.1619 < href <= 47.3501
#         W = 1 + ((hp - 39.7499) / 89.4107)
#         rho = 0.0032618 * rho0 * W^(-13.2011)
#         T = 250.350 * W
#         Pa = 0.0028338 * Pa0_atmos * W^(-12.2011)
#     elseif 47.3501 < href <= 51.4125
#         W = exp((48.6252 - hp) / 7.9223)
#         rho = 0.00094920 * rho0 * W
#         T = 270.650
#         Pa = 8.9155e-4 * W
#     elseif 51.4125 < href <= 71.8020
#         W = 1 - ((hp - 59.4390) / 88.2218)
#         rho = 0.00025280 * rho0 * W^11.2011
#         T = 247.021 * W
#         Pa = 2.1671e-4 * Pa0_atmos * W^12.2011
#     elseif 71.8020 < href <= 86.000
#         W = 1 - ((hp - 78.0303) / 100.2950)
#         rho = 0.000017632 * rho0 * W^16.0816
#         T = 200.590 * W
#         Pa = 1.2274e-5 * Pa0_atmos * W^17.0816
#     elseif 86.000 < href <= 91.000
#         W = exp((87.2848 - hp) / 5.4700)
#         rho = 0.0000036411 * rho0 * W
#         T = 186.8700
#         Pa = (2.2730 + 1.042e-3 * hp) * 1e-6 * W
#     else
#         W = -3.41173e-6 * href^3 + 0.0018157 * href^2 - 0.337816 * href + 12.63404
#         rho = 10.0^W
#         T = 186.8700
#         Pa = 0
#     end
#     rho =rho0*exp(-href/7.110);
#     soundv = 20.0468 * sqrt(T)
#     return rho, soundv
# end
function m_GetDensity(href)
    #大气密度分段
    #h单位是km
    R = 6356.766
    rho0 = 1.225
    href = clamp(href, 0.0, 120.0)
    hp = R * href / (R + href)
    if href <= 11.0191
        W = 1 - (hp / 44.3308)
        rho = rho0 * W^4.2559
        T = 288.15 * W
        Pa = Pa0_atmos * W^5.2559
    elseif 11.0191 < href <= 20.0631
        W = exp((14.9647 - hp) / 6.3416)
        rho = 0.15898 * rho0 * W
        T = 216.650
        Pa = Pa0_atmos * 0.11953 * W
    elseif 20.0631 < href <= 32.1619
        W = 1 + ((hp - 24.9021) / 221.552)
        rho = 0.032722 * rho0 * W^(-35.1629)
        T = 221.552 * W
        Pa = 0.025158 * Pa0_atmos * W^(-34.1629)
    elseif 32.1619 < href <= 47.3501
        W = 1 + ((hp - 39.7499) / 89.4107)
        rho = 0.0032618 * rho0 * W^(-13.2011)
        T = 250.350 * W
        Pa = 0.0028338 * Pa0_atmos * W^(-12.2011)
    elseif 47.3501 < href <= 51.4125
        W = exp((48.6252 - hp) / 7.9223)
        rho = 0.00094920 * rho0 * W
        T = 270.650
        Pa = 8.9155e-4 * W
    elseif 51.4125 < href <= 71.8020
        W = 1 - ((hp - 59.4390) / 88.2218)
        rho = 0.00025280 * rho0 * W^11.2011
        T = 247.021 * W
        Pa = 2.1671e-4 * Pa0_atmos * W^12.2011
    elseif 71.8020 < href <= 86.000
        W = 1 - ((hp - 78.0303) / 100.2950)
        rho = 0.000017632 * rho0 * W^16.0816
        T = 200.590 * W
        Pa = 1.2274e-5 * Pa0_atmos * W^17.0816
    elseif 86.000 < href <= 91.000
        W = exp((87.2848 - hp) / 5.4700)
        rho = 0.0000036411 * rho0 * W
        T = 186.8700
        Pa = (2.2730 + 1.042e-3 * hp) * 1e-6 * W
    else
        W = -3.41173e-6 * href^3 + 0.0018157 * href^2 - 0.337816 * href + 12.63404
        rho = 10.0^W
        T = 186.8700
        Pa = 0
    end

    soundv = 20.0468 * sqrt(T)
    return rho, soundv
end
#########实际的常微分方程
function trueODEfunc(du, u, p, t)
    #再入动力学方程
    (r, theta, phi, V, gamma, psi, alpha, sigma) = u
    dalpha = 0.02 * p[1] * (V - 2000) / 1E3
    dsigma = 0.02 * p[2] * (r - R0) / 1E6
    (dalpha, dsigma) = mybound([dalpha; dsigma])#对网络输出量化到指定区间内
    alpha = clamp(alpha, minalpha, magalpha)#角度限幅
    sigma = clamp(sigma, -magsigma, magsigma)

    height = r - R0
    rou, soundV = m_GetDensity(height / 1000.0)

    Ma = V / soundV
    q = 0.5 * rou * V * V#动压
    dQ = kq * ((rou / rho0)^0.5) * (V / V_scale)^3.15#热流

    alpha = alpha * RTD#公式里单位是角度,攻角应为大于等于0的数为合理
    #气动力/质量
    D = q * Sref / (mass) * (ad + bd * Ma + cd * alpha + dd * Ma * Ma + ed * Ma * alpha + fd * alpha * alpha)
    L = q * Sref / (mass) * (al + bl * Ma + cl * alpha + dl * Ma * Ma + el * Ma * alpha + fl * alpha * alpha)
    alpha = alpha * DTR

    nn = sqrt(D * D + L * L)#过载

    g = miu / (r * r)
    #考虑地球自转角速度的影响
    du .= [V * sin(gamma);
        V * cos(gamma) * sin(psi) / (r * cos(phi));
        V * cos(gamma) * cos(psi) / r;
        -D - g * sin(gamma) + OMEGA * OMEGA * r * cos(phi) * (sin(gamma) * cos(phi) - cos(gamma) * sin(phi) * cos(psi));
        (1 / V) * (L * cos(sigma) - (g - V * V / r) * cos(gamma)) + 2 * OMEGA * cos(phi) * sin(psi) + OMEGA * OMEGA * r * cos(phi) / V * (cos(gamma) * cos(phi) + sin(gamma) * sin(phi) * cos(psi));
        (1 / V) * (L * sin(sigma) / cos(gamma)) + (V / r) * cos(gamma) * sin(psi) * tan(phi) + 2 * OMEGA * (sin(phi) - cos(phi) * cos(psi) * tan(gamma)) + OMEGA * OMEGA * r * sin(psi) * sin(phi) * cos(phi) / V / cos(gamma);
        dalpha;
        dsigma]
end

condition(u, t, integrator) = u[1] <= R0
affect!(integrator) = terminate!(integrator)
ODEcallback = DiscreteCallback(condition, affect!)

probtrue = ODEProblem(trueODEfunc, zeros(statenum), p=zeros(), tspan);

if sampsign#重新生成样本
    ps = Float32.(2 * rand(sampnum, 2) .- 1)
    samples = Array{Any}(undef, sampnum, 2)
    for ii in 1:sampnum
        r = (55e3 + 20e3 * rand()) + R0
        theta = (-10 + 20 * rand()) * DTR
        phi = (-10 + 20 * rand()) * DTR
        V = (5.0e3 + 2.0e3 * rand())
        gamma = -(2 * rand()) * DTR
        psi = (-180 + 360 * rand()) * DTR
        alpha = (15 + 5 * rand()) * DTR
        sigma = (-0.1 + 0.2 * rand()) * DTR
        samples[ii, 1] = [r; theta; phi; V; gamma; psi; alpha; sigma]#初始状态
    end

    function generate_probstrue(samples, ps)
        return function (prob, i, repeat)
            remake(prob, u0=samples[i, 1], p=ps[i, :])
        end
    end
    ensemble_probtrue = EnsembleProblem(probtrue, prob_func=generate_probstrue(samples, ps))
    sol_full = solve(ensemble_probtrue, RK4(), dt=0.1, myEnsemble, saveat=tstepsdata, trajectories=sampnum)#, callback = ODEcallback
    us = zeros(statenum, 0)#状态量样本
    for ii in 1:sampnum
        global us
        u = reduce(hcat, sol_full[ii].u)
        us = hcat(us, u)
        (r, theta, phi, V, gamma, psi, alpha, sigma) = Array(sol_full[ii].u)[end]#终端状态    
        samples[ii, 2] = [theta; phi; r; V; 50e3; gamma]#以终端状态作为目标
    end
    save("sample.jld2", "us", us, "samples", samples)#存样本
else#加载样本
    us = load("sample.jld2", "us")#加载模型权重
    samples = load("sample.jld2", "samples")#加载模型内部状态
end
#计算网络输入的归一化量
meanu = mean(us, dims=2);
varu = var(us, dims=2);

goals = reduce(hcat, samples[:, 2]);
meang = mean(goals, dims=2);
varg = var(goals, dims=2);

allmean = [meanu; meang];
allvar = [varu; varg];
#########神经网络与包含的常微分方程
# 初始化模型
policynet = PolicyNet19(statenum + goalnum, 128, actnum)#隐含层单元数量
p_NN, st = Lux.setup(rng, policynet);
#给网络bn赋值，在训练中也不变
newbn = setproperties(st.bn, (running_mean=Float32.(reshape(allmean, statenum + goalnum)), running_var=Float32.(reshape(allvar, statenum + goalnum))));
st = setproperties(st, (bn = newbn));
#model能运行测试
# st=Lux.initialstates(rng,policynet)
inp = 20 * randn(statenum + goalnum)

u, stt = policynet(inp, p_NN, st)
println("Model output: ", u)

function neuralODEfunc(du, u, p, t)
    #再入动力学方程
    (r, theta, phi, V, gamma, psi, alpha, sigma, thetaT, phiT, rf, Vf, Sf, gammaf) = u

    nety, stt = policynet(u[1:statenum+goalnum], p, st)
    (dalpha, dsigma) = mybound(nety)#对网络输出量化到指定区间内
    # if r > R0
    alpha = clamp(alpha, minalpha, magalpha)#角度限幅
    sigma = clamp(sigma, -magsigma, magsigma)

    height = r - R0
    rou, soundV = m_GetDensity(height / 1000.0)

    Ma = V / soundV
    q = 0.5 * rou * V * V#动压
    dQ = kq * ((rou / rho0)^0.5) * (V / V_scale)^3.15#热流

    alpha = alpha * RTD#公式里单位是角度,攻角应为大于等于0的数为合理
    #气动力/质量
    D = q * Sref / (mass) * (ad + bd * Ma + cd * alpha + dd * Ma * Ma + ed * Ma * alpha + fd * alpha * alpha)
    L = q * Sref / (mass) * (al + bl * Ma + cl * alpha + dl * Ma * Ma + el * Ma * alpha + fl * alpha * alpha)
    alpha = alpha * DTR

    nn = sqrt(D * D + L * L)#过载

    g = miu / (r * r)
    #考虑地球自转角速度的影响

    # typeu = eltype(u)
    du .= [
        V * sin(gamma);
        V * cos(gamma) * sin(psi) / (r * cos(phi));
        V * cos(gamma) * cos(psi) / r;
        -D - g * sin(gamma) + OMEGA * OMEGA * r * cos(phi) * (sin(gamma) * cos(phi) - cos(gamma) * sin(phi) * cos(psi));
        (1 / V) * (L * cos(sigma) - (g - V * V / r) * cos(gamma)) + 2 * OMEGA * cos(phi) * sin(psi) + OMEGA * OMEGA * r * cos(phi) / V * (cos(gamma) * cos(phi) + sin(gamma) * sin(phi) * cos(psi));
        (1 / V) * (L * sin(sigma) / cos(gamma)) + (V / r) * cos(gamma) * sin(psi) * tan(phi) + 2 * OMEGA * (sin(phi) - cos(phi) * cos(psi) * tan(gamma)) + OMEGA * OMEGA * r * sin(psi) * sin(phi) * cos(phi) / V / cos(gamma);
        dalpha;
        dsigma;
        0.0 * dalpha;
        0.0 * dalpha;
        0.0 * dalpha;
        0.0 * dalpha;
        0.0 * dalpha;
        0.0 * dalpha]
    # else
    #     du .= zeros(statenum+goalnum+pathnum)
    # end
end

function genloss(goal, state)
    #计算损失值
    (thetaT, phiT, rf, Vf, Sf, gammaf) = goal#目标状态
    (r, theta, phi, V, gamma, psi, alpha, sigma) = state#实际终端状态
    S = R0 * acos(sin(phi) * sin(phiT) + cos(phi) * cos(phiT) * cos(thetaT - theta))#剩余航程
    Sloss = max(S - Sf, 0)#这里是要求目标状态到实际终端状态航程不大于某个数
    Vloss = max(abs(V - Vf) - Verrormax, 0)
    rloss = max((abs(r - rf) - rerrormax), 0)
    loss = [Sloss, Vloss, rloss]
    return loss
end

function genALLloss(goal, state)
    #计算损失值
    (thetaT, phiT, rf, Vf, Sf, gammaf) = goal#目标状态
    (r, theta, phi, V, gamma, psi, alpha, sigma) = state#实际终端状态
    S = R0 * acos(sin(phi) * sin(phiT) + cos(phi) * cos(phiT) * cos(thetaT - theta))#剩余航程
    Sloss = max(S - Sf, 0)#这里是要求目标状态到实际终端状态航程不大于某个数
    Vloss = max(abs(V - Vf) - Verrormax, 0)
    rloss = max((abs(r - rf) - rerrormax), 0)
    return Sloss, Vloss, rloss
end

nu = Float32[2.0; 0.0];
nu0 = Float32[nu; 0.5];
prob = ODEProblem(neuralODEfunc, nu0, tspan);

# DP8(), reltol=1e-4, abstol=1e-1 RK4(), dt=ndt 
#Tsit5(), abstol=1e-1, reltol=1e-4
# Vern9(), abstol=1e-1, reltol=1e-5
function loss_neuralode(p_NN, extra_params::NamedTuple)
    samples = extra_params.samples
    lnorm = extra_params.lnorm
    sampnum = size(samples, 1)
    loss = 0.0
    for ii in 1:sampnum
        loss = loss + sum(genloss(samples[ii, 2], (Array(solve(remake(prob; u0=[samples[ii, 1]; samples[ii, 2]]), RK4(), dt=ndt; p=p_NN, sensealg=InterpolatingAdjoint(autojacvec=ZygoteVJP()), saveat=tsteps)))[1:statenum, end]) ./ lnorm)
    end
    loss = loss / sampnum
    return loss
end
##############用于画图的callback
function Xloss_neuralode(p_NN, samples)
    #有中间结果
    sampnum = size(samples, 1)
    losss = zeros(Float32, sampnum, 3)
    preds = Array{Any}(undef, sampnum)#Julia中循环内部声明的变量不会存下来
    for ii in 1:sampnum
        preds[ii] = Array(solve(remake(prob; u0=[samples[ii, 1]; samples[ii, 2]]), RK4(), dt=0.1; p=p_NN, saveat=tsteps));#, callback=ODEcallback
        (losss[ii, 1], losss[ii, 2], losss[ii, 3]) = genALLloss(samples[ii, 2], (preds[ii])[1:statenum, end])
    end
    loss = mean(losss, dims=1)
    return loss, preds, losss
end

callback = function (l, preds, losss)
    # push!(loss_values, l)
    open("losshis.txt", "a") do io
        println(io, l)
    end
    open("alllosshis.txt", "a") do io
        println(io, losss)#存所有的loss
    end
    println(l)
    #sdloss = mean(losss, dims=1)#均值
    sdloss = mean(abs.(losss), dims=1)#绝对值的均值
    return sdloss#自定义的损失函数模
end
##############训练过程中用于存损失函数的callback
open("losshis.txt", "w") do io#清空之前的loss文件
    write(io, "")
end
open("alllosshis.txt", "w") do io#清空之前的loss文件
    write(io, "")
end

loss_values = Float32[];
iter = 0;

newloss = callback(Xloss_neuralode(ComponentArray(p_NN), samples)...);
lnorm = max.(newloss, stdeps);
##声明优化问题并求解
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((p, extra_params) -> loss_neuralode(p, extra_params), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentArray(p_NN), (samples=samples, lnorm=lnorm))

# random_indices = random_grouping(sampnum, batchnum);

for i = 1:expochs
    global p_NN, lnorm
    random_indices = rand(1:sampnum, batchnum)
    selected_samples = samples[random_indices, :]#取一些数来训练
    mynorm = lnorm#归一化用均值
    newprob = remake(optprob; u0=ComponentArray(p_NN), p=(samples=selected_samples, lnorm=mynorm))
    result_neuralode = Optimization.solve(newprob, OptimizationOptimisers.OptimiserChain(OptimizationOptimisers.ClipNorm(10), OptimizationOptimisers.RAdam(lr));
        maxiters=(inpochs + 1))#定次数优化时，优化(maxiter-1)次
    # result_neuralode = DiffEqFlux.sciml_train(newprob, OptimizationOptimisers.RAdam(lr), maxiters=2, callback=ODEcallback)#定次数优化时，优化(maxiter-1)次
    p_NN = result_neuralode.u
    newprob = remake(newprob; u0=ComponentArray(p_NN))
    # if i % BFGSpoch == 0
    #     # result_neuralode = Optimization.solve(remake(optprob; u0=ComponentArray(p_NN), p=(samples=selected_samples,)), Optim.BFGS(; initial_stepnorm=0.002); maxiters=(inBFGSpochs + 1), allow_f_increases=false)
    #     result_neuralode = Optimization.solve(newprob, OptimizationOptimisers.OptimiserChain(OptimizationOptimisers.ClipNorm(10), OptimizationOptimisers.AdaBelief(lr * 10.0));
    #         maxiters=(inBFGSpochs + 1))#定次数优化时，优化(maxiter-1)次
    #     p_NN = result_neuralode.u
    # end
    if i % cpochs == 0
        save("model$i.jld2", "model_structure", policynet, "model_weight", p_NN, "model_states", st)#存模型

        newloss = callback(Xloss_neuralode(p_NN, samples)...)
        lnorm = max.(ngreedy * newloss + (1 - ngreedy) * lnorm, stdeps)
    end
end
########画损失函数曲线
# plt = plot(loss_values, label="Loss")
# title!(plt, "Training Loss Over Time")# 添加标题和标签
# xlabel!(plt, "Iteration")
# ylabel!(plt, "Loss Value")
# plot!(plt, grid=true)# 添加网格线
# display(plt)
# savefig(plt, "loss.pdf")

# policynet = load("model.jld2", "model_structure")#加载模型结构
# p_NN = load("model.jld2", "model_weight")#加载模型权重
# st = load("model.jld2", "model_states")#加载模型内部状态